import torch
from torch import nn

from .build import ModelFactory
from ..model.efficientdet import EfficientDetBackbone
from ..loss import FocalLoss, FocalCIoULoss
from utils.box import box_transform


@ModelFactory.register()
class EfficientDet(nn.Module):
    def __init__(self, model_config):
        super(EfficientDet, self).__init__()
        self.criterion = FocalLoss(
            classification_loss_weight=model_config['loss']['classification_loss_weight'],
            regression_loss_weight=model_config['loss']['regression_loss_weight'],
        )
        self.model = EfficientDetBackbone.from_pretrained(
            num_classes=model_config['num_classes'],
            compound_coef=model_config['compound_coef'],
            load_weights=model_config['load_weight'],
            in_channels=model_config['in_channels'],
            ratios=model_config['ratios'],
            scales=model_config['scales'],
        )

    def forward(self, data):
        images = data['images']
        annotation_boxes = data.get('boxes', None)
        annotation_category_ids = data.get('category_ids', None)

        regression, classification, anchors = self.model(images)

        boxes = box_transform(anchors, regression)
        loss_dict = {}
        if annotation_boxes is not None and annotation_category_ids is not None:
            classification_loss, regression_loss = self.criterion(classification, regression, anchors, annotation_boxes, annotation_category_ids)
            loss_dict['classification_loss'] = classification_loss
            loss_dict['regression_loss'] = regression_loss
        return loss_dict, {"boxes": boxes, "classification": classification}


@ModelFactory.register()
class EfficientDet_FocalCIoULoss(nn.Module):
    def __init__(self, model_config):
        super(EfficientDet, self).__init__()
        self.criterion = FocalCIoULoss(
            classification_loss_weight=model_config['loss']['classification_loss_weight'],
            regression_loss_weight=model_config['loss']['regression_loss_weight'],
        )
        self.model = EfficientDetBackbone.from_pretrained(
            num_classes=model_config['num_classes'],
            compound_coef=model_config['compound_coef'],
            load_weights=model_config['load_weight'],
            in_channels=model_config['in_channels'],
            ratios=model_config['ratios'],
            scales=model_config['scales'],
        )

    def forward(self, data):
        images = data['images']
        annotation_boxes = data.get('boxes', None)
        annotation_category_ids = data.get('category_ids', None)

        regression, classification, anchors = self.model(images)

        boxes = box_transform(anchors, regression)
        loss_dict = {}
        if annotation_boxes is not None and annotation_category_ids is not None:
            classification_loss, regression_loss = self.criterion(classification, regression, anchors, annotation_boxes, annotation_category_ids)
            loss_dict['classification_loss'] = classification_loss
            loss_dict['regression_loss'] = regression_loss
        return loss_dict, {"boxes": boxes, "classification": classification}
